/** 
 * @file   UpgmaAlgorithm.cpp
 * @author Scanagatta Mauro - Alberto Franco
 * 
 * Contains the implementation of the UPGMA Algorithm we did as first part of 
 * the assignment.
 */
#include <cstdlib>
#include <iostream>
#include <math.h>
#include <stdio.h>

#include "Align.h"
#include "ScoringS2S.h"
#include "SequenceData.h"
#include "NWAlign.h"
#include "AGPFunction.h"

#include "UpgmaAlgorithm.h"
#include "NewickTreeFactory.h"

namespace Biopool
{

#if !defined(max)
#define max(a, b) ((a) > (b) ? (a) : (b))
#define min(a, b) ((a) < (b) ? (a) : (b))
#endif

/**
 * Contructor. Takes the input sequences and create the data structures needed
 * to solve the probelm.
 * @param inputSequences The sequences to build up the evolutionary tree
 */
UpgmaAlgorithm::UpgmaAlgorithm(std::vector<std::string> & inputSequences) :
    sequences(inputSequences) 
{
	/// Set distances parameters
    this->distanceGapPenalty = 10;
    this->distanceGapExtension = 2;	
	/// Create Align classes 
    std::ifstream matrixFile("data/blosum62.dat");
	if (!matrixFile) 
        throw "Error opening structure matrix file.";
        
    this->subMatrix = new SubMatrix(matrixFile);
    this->gapFunction = new AGPFunction(this->distanceGapPenalty, this->distanceGapExtension);
}

/// Copy constructor
UpgmaAlgorithm::UpgmaAlgorithm(const UpgmaAlgorithm& orig) :
sequences(orig.sequences) { }

/// Deconstructor
UpgmaAlgorithm::~UpgmaAlgorithm() 
{
	delete(this->subMatrix);
	delete(this->gapFunction);
}


std::string UpgmaAlgorithm::execute() 
{
    std::cout << "\t \t \t * Starting Upgma Algorithm... * " 
              << std::endl << std::endl;
    /// Inizialization
    std::cout << "\t \t \t * Calcolo delle distanze iniziali... * " 
              << std::endl << std::endl;
    /// Calculate distance matrix
    this->calculateInitialDistanceMatrix();
    /// Create initial clusters
    for (unsigned int i = 0; i < sequences.size(); i++) 
    {
        intSet * setTemp = new intSet;
        setTemp->insert(i);
        NewickTree::Node * temp = new NewickTree::Node(setTemp);
        this->clusters.push_back(temp);
    }

    int numClusters = this->clusters.size();
    NewickTree::Node * nodeRoot = 0;

    std::cout << "\t \t \t * Joining the clusters... * " << std::endl << std::endl;
    
    while (numClusters > 1) 
    {

        /// a, b: Cluster indexes to pair in this iteration
        /// l: Distance between clusters
        int a, b;
        float l;
        /// Find minimum L(a,b)
        this->getMinimumDistance(a, b, l);

        std::cout << "Unione: a" << a << ", b: " << b << ", ed l: " << l << "\n" << std::endl;

        this->updateDistance(a, b);

        nodeRoot = this->updateClusters(a, b, l);

        /// Update clusters count
        numClusters--;
    }
    
    std::cout << "\t \t \t * Tree built! * " << std::endl << std::endl;
    NewickTree newickTree(nodeRoot);
    
    std::cout << "\t \t \t * Tree Serialization...  * " << std::endl << std::endl;
    return NewickTreeFactory::serialize(newickTree);
}

/// Distance calculation
float UpgmaAlgorithm::calculateDistance(std::string firstSeq, std::string secondSeq) 
{
    this->sequenceData = new SequenceData(2, firstSeq, secondSeq);
    this->scoringScheme = new ScoringS2S(this->subMatrix, this->sequenceData, 0, 1.00);
    this->align = new NWAlign(sequenceData, this->gapFunction, this->scoringScheme);

	this->align->doMatchScore(cout);
	
    /// Get alignment
    vector<string> match = this->align->getMatch(); 
    cout << std::endl << "Match: " << std::endl;
	for (unsigned int i = 0; i < match.size(); ++i)
        cout << match[i] << std::endl;
        
    firstSeq = match[0];
    secondSeq = match[1];
    
    /// Compute different residues
    int diffResidues = 0;
    int n = min(firstSeq.size(), secondSeq.size());

    for (int i = 0; i < n; i++) 
    {
        if (firstSeq[i] != secondSeq[i])
            diffResidues++;
    }

    float residuesFract = ((float) diffResidues) / n;

    if (residuesFract > 0.7499)
        residuesFract = 0.7499;
        
    delete(this->sequenceData);
    delete(this->scoringScheme);
    delete(this->align);

    return residuesFract;
}

/// Return the distances index.
int UpgmaAlgorithm::getDistanceIndex(int firstId, int secondId) 
{
    int n = this->sequences.size();

    if (firstId > n - 1 || firstId < 0)
        throw "distanceMatrix access with wrong firstID";

    if (secondId > n - 1 || secondId < 0)
        throw "distanceMatrix access with wrong secondID";

    if (firstId == secondId)
        return 0;

    int i = max(firstId, secondId);
    int j = min(firstId, secondId);

    return (j + ((i - 1) * i) / 2);
}

/// Calculate initial distance matrix.
void UpgmaAlgorithm::calculateInitialDistanceMatrix() 
{
    std::cout << "Calculating Initial Distance Matrix... " << std::endl;

    int n = this->sequences.size();
    /// Dimension of distance matrix, optimized for using only upper 
    /// matrix without the diagonal.
    int m = ((n * n) - n) / 2;

    this->distanceMatrix.resize(m);

    /// Distance matrix calculation (upper only without diagonal)
    for (int i = 0; i < n; i++) 
    {
        for (int j = 0; j < i; j++) 
        {
            /// Calculate distance between sequence i and j
            float residuesFract = this->calculateDistance(
                    this->sequences[i], this->sequences[j]);
            float distance = -3.0 / 4 * log(1 - 4.0 * residuesFract / 3);
            printf("   -> distance %d e %d : %4.2f \n ", i, j, distance);

            /// Insert value into the matrix, with efficent allocation
            int index = this->getDistanceIndex(i, j);
            this->distanceMatrix[index] = distance;
        }
    }
}


/// Search for a and b, clusters with minimum distance
void UpgmaAlgorithm::getMinimumDistance(int & a, int & b, float & l) 
{
    l = 100000;
    int n = this->sequences.size();

    for (int i = 0; i < n; i++)
        for (int j = 0; j < i; j++) 
        {
            int index = this->getDistanceIndex(i, j);
            float distance = this->distanceMatrix[index];
            // If a better distance is discovered
            if (distance > 0 && distance < l) 
            {
                a = i;
                b = j;
                l = distance;
            }
        }

    if (!l > 0) 
    {
        throw "No non-zero distance found.";
    }
}

/// Update the clusters with the distance
NewickTree::Node * UpgmaAlgorithm::updateClusters(int a, int b, float l) 
{
    /// Take a and b clusters, create a new node that contains both
    NewickTree::Node * nodeA = this->clusters[a];
    NewickTree::Node * nodeB = this->clusters[b];

    /// Add to K's set all sequences in a and b clusters
    intSet * seqK = new intSet;
    intSet::iterator setIt;
    intSet * seqA = nodeA->getSequencesSet();
    for (setIt = seqA->begin(); setIt != seqA->end(); setIt++)
        seqK->insert(*setIt);
    intSet * seqB = nodeB->getSequencesSet();
    for (setIt = seqB->begin(); setIt != seqB->end(); setIt++)
        seqK->insert(*setIt);

    /// Create the new cluster, l/2 high, with K's set and a and b as childs.
    NewickTree::Node * nodeK = new NewickTree::Node(l / 2, seqK, nodeA, nodeB);

    /// Update node vector
    this->clusters[a] = nodeK;
    this->clusters[b] = 0;

    return nodeK;
}

/**
 * Update all costs in the distance table.
 * @param a Index of all counts toward (a, b)
 * @param b Index to be set at -1 
 */
void UpgmaAlgorithm::updateDistance(int a, int b) 
{
    int sizeA = this->clusters[a]->getSequencesSet()->size();
    int sizeB = this->clusters[b]->getSequencesSet()->size();

    int n = this->sequences.size();

    /// Updata costs: new cluster A
    for (int j = 0; j < a; j++) 
    {
        int indexAL = this->getDistanceIndex(a, j);
        int indexBL = this->getDistanceIndex(b, j);
        this->distanceMatrix[indexAL] = 
                ((this->distanceMatrix[indexAL] * sizeA +
                this->distanceMatrix[indexBL] * sizeB) / (sizeA + sizeB));
    }

    for (int i = a; i < n; i++) 
    {
        int indexAL = this->getDistanceIndex(a, i);
        int indexBL = this->getDistanceIndex(b, i);
        this->distanceMatrix[indexAL] = 
            ((this->distanceMatrix[indexAL] * sizeA + 
              this->distanceMatrix[indexBL] * sizeB ) / (sizeA + sizeB));
    }

    /// Update costs: old cluster B
    for (int j = 0; j < b; j++) 
    {
        int index = this->getDistanceIndex(b, j);
        this->distanceMatrix[index] = -1;
    }

    for (int i = b; i < n; i++) 
    {
        int index = this->getDistanceIndex(b, i);
        this->distanceMatrix[index] = -1;
    }

}

} /* namespace */
